import numpy as np
from utils.LogUtil import my_logger
from batchgenerators.transforms.spatial_transforms import SpatialTransform

"""
现在用的Batchgenerators和pytorch兼容并不是太好。
后续可以考虑使用https://github.com/ncullen93/torchsample来做数据增强
"""


class CustomPreprocessModule(object):
    def __init__(self, patch_shape, probability=0.3):
        self.data_key = "train_data"
        self.label_key = "label"
        self.probability = probability
        self.transform = SpatialTransform(patch_shape, np.array(patch_shape) // 2,
                                          do_elastic_deform=False, alpha=(0., 1500.), sigma=(30., 50.),
                                          do_rotation=True, angle_z=(0, 2 * np.pi),
                                          do_scale=True, scale=(0.8, 1.2),
                                          border_mode_data='constant', border_cval_data=0, order_data=1,
                                          data_key=self.data_key,
                                          label_key=self.label_key,
                                          random_crop=False)

    def process(self, train_data, label):
        random_number = np.random.rand()
        if random_number < self.probability:
            result_dict = self.transform.__call__(train_data=train_data, label=label)
            my_logger.info("Once data augmentation has been completed!")
            return result_dict.get(self.data_key), result_dict.get(self.label_key)
        else:
            return train_data, label
